import os
from random import choice
import shutil

classes =['Abyssinian', 'Bengal', 'Birman', 'Bombay', 'British_Shorthair', 'Egyptian_Mau', 'Maine_Coon', 'Persian', 'Ragdoll', 'Russian_Blue', 'Siamese', 'Sphynx']

# arrays to store file names
imgs = [[] for i in range(len(classes))]
xmls = [[] for j in range(len(classes))]

# setup dir names
# trainPath = 'D:/yolov5_traffic_sign_detection/dataset/images/train'
# valPath = 'D:/yolov5_traffic_sign_detection/dataset/images/val'
crsPath = 'D:/work/images'  # dir where images and annotations stored

# setup ratio (val ratio = rest of the files in origin dir after splitting into train and test)
train_ratio = 0.7
val_ratio = 0.15
test_ratio = 0.15
# total count of imgs
totalImgCount = len(os.listdir(crsPath)) / 2


def findThename(filename):
    for i in range(len(classes)):
        if classes[i] in filename:
            return i


# soring files to corresponding arrays
for (dirname, dirs, files) in os.walk(crsPath):
    for filename in files:
        if filename.endswith('.txt'):
            xmls[findThename(filename)].append(filename)
        else:
            imgs[findThename(filename)].append(filename)

# counting range for cycles
countForTrain = [int(len(i) * train_ratio) for i in imgs]
countForVal = [int(len(i) * val_ratio) for i in imgs]
countForTest = [int(len(i) * test_ratio) for i in imgs]
print("training images are : ", countForTrain)
print("Validation images are : ", countForVal)
print("Testidation images are : ", countForTest)

trainimagePath = 'D:/work/dataset/images/train'
trainlabelPath = 'D:/work/dataset//labels/train'
valimagePath = 'D:/work/dataset/images/val'
vallabelPath = 'D:/work/dataset//labels/val'
testimagePath = 'D:/work/dataset/images/test'
testlabelPath = 'D:/work/dataset//labels/test'
# cycle for val dir
for x in range(len(countForVal)):
    for j in range(countForVal[x]):
        fileJpg = choice(imgs[x])  # get name of random image from origin dir
        fileXml = fileJpg[:-4] + '.txt'  # get name of corresponding annotation file

        # move both files into train dir
        # shutil.move(os.path.join(crsPath, fileJpg), os.path.join(trainimagePath, fileJpg))
        # shutil.move(os.path.join(crsPath, fileXml), os.path.join(trainlabelPath, fileXml))
        shutil.copy(os.path.join(crsPath, fileJpg), os.path.join(valimagePath, fileJpg))
        shutil.copy(os.path.join(crsPath, fileXml), os.path.join(vallabelPath, fileXml))

        # remove files from arrays
        imgs[x].remove(fileJpg)
        xmls[x].remove(fileXml)
# cycle for test dir
for x in range(len(countForTest)):
    for j in range(countForTest[x]):
        fileJpg = choice(imgs[x])  # get name of random image from origin dir
        fileXml = fileJpg[:-4] + '.txt'  # get name of corresponding annotation file

        # move both files into train dir
        # shutil.move(os.path.join(crsPath, fileJpg), os.path.join(valimagePath, fileJpg))
        # shutil.move(os.path.join(crsPath, fileXml), os.path.join(vallabelPath, fileXml))
        shutil.copy(os.path.join(crsPath, fileJpg), os.path.join(testimagePath, fileJpg))
        shutil.copy(os.path.join(crsPath, fileXml), os.path.join(testlabelPath, fileXml))

        # remove files from arrays
        imgs[x].remove(fileJpg)
        xmls[x].remove(fileXml)

# cycle for train dir
for x in range(len(countForTrain)):
    for j in range(countForTrain[x]):
        fileJpg = choice(imgs[x])  # get name of random image from origin dir
        fileXml = fileJpg[:-4] + '.txt'  # get name of corresponding annotation file

        # move both files into train dir
        # shutil.move(os.path.join(crsPath, fileJpg), os.path.join(valimagePath, fileJpg))
        # shutil.move(os.path.join(crsPath, fileXml), os.path.join(vallabelPath, fileXml))
        shutil.copy(os.path.join(crsPath, fileJpg), os.path.join(trainimagePath, fileJpg))
        shutil.copy(os.path.join(crsPath, fileXml), os.path.join(trainlabelPath, fileXml))

        # remove files from arrays
        imgs[x].remove(fileJpg)
        xmls[x].remove(fileXml)


# rest of files will be validation files, so rename origin dir to val dir
# os.rename(crsPath, valPath)
# shutil.move(crsPath, valPath)
